package org.apache.flink.streaming.api.datastream;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.metrics.Counter;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.HashMap;
import java.util.TreeMap;

@SuppressWarnings("unused")
public abstract class LocalMergeFunction<K, IN, OUT> extends RichFlatMapFunction<IN, OUT> {
    long emitTime;
    int emitCount;

    Counter total, merged;

    HashMap<K, TreeMap<Long, OUT>> map;
    long lastEmit;
    int count;

    public LocalMergeFunction(Duration emitTime) {
        this.emitTime = emitTime.toMillis();
    }

    public LocalMergeFunction(int emitCount) {
        this.emitCount = emitCount;
    }

    @Override
    public void open(Configuration parameters) {
        map = new HashMap<>();
        lastEmit = System.currentTimeMillis();
        count = 0;
        total = getRuntimeContext().getMetricGroup().counter("localMergeTotal");
        merged = getRuntimeContext().getMetricGroup().counter("localMergeMerged");
    }

    @Override
    public void flatMap(IN in, Collector<OUT> out) {
        merge(in);
        total.inc();
        count++;
        if (shouldEmit(in) || shouldEmit()) {
            emit(out);
        }
    }

    public abstract K getKey(IN in);
    public abstract long getSlot(IN in);
    public abstract OUT merge(IN in, OUT old);
    public abstract boolean shouldEmit(IN in);

    private void merge(IN in) {
        map.putIfAbsent(getKey(in), new TreeMap<>());
        map.get(getKey(in)).compute(getSlot(in), (k, v) -> {
            if (v != null) {
                merged.inc();
            }
            return merge(in, v);
        });
    }

    private void emit(Collector<OUT> out) {
        map.values().forEach(v -> {
            v.values().forEach(out::collect);
            v.clear();
        });
    }

    private boolean shouldEmit() {
        if (emitTime > 0) {
            long now = System.currentTimeMillis();
            long diff = now - lastEmit;
            if (diff > emitTime) {
                emitTime = now;
                return true;
            }
            return false;
        }
        if (emitCount > 0) {
            if (count >= emitCount) {
                count = 0;
                return true;
            }
            return false;
        }
        return false;
    }
}
